import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

from matplotlib.patches import Rectangle


def plot_double_axis_with_background(df, col1, col2, title="双序列趋势对比图"):
    """
    绘制双轴折线图，根据同趋势数据设置背景颜色（修复背景不显示问题）
    
    参数:
    df: 包含数据的DataFrame，需包含col1、col2和"同趋势"列，索引为日期
    col1: 左轴序列的列名
    col2: 右轴序列的列名
    title: 图表标题
    """
    # 1. 基础配置（解决中文和负号显示）
    plt.rcParams["font.family"] = ["SimHei"]
    plt.rcParams["axes.unicode_minus"] = False
    fig, ax1 = plt.subplots(figsize=(12, 6))

    # 2. 先绘制所有数据（确保轴范围已根据数据调整）
    # 左轴：col1（蓝色实线+圆形标记）
    color1 = 'tab:blue'
    ax1.set_xlabel('日期', fontsize=10)
    ax1.set_ylabel(col1, color=color1, fontsize=10)
    line1 = ax1.plot(df.index, df[col1], color=color1, marker='o', linestyle='-', 
                     linewidth=2, markersize=6, label=col1)
    ax1.tick_params(axis='y', labelcolor=color1)
    ax1.legend(loc='upper left', fontsize=9)

    # 右轴：col2（红色虚线+方形标记）
    ax2 = ax1.twinx()  
    color2 = 'tab:red'
    ax2.set_ylabel(col2, color=color2, fontsize=10)  
    line2 = ax2.plot(df.index, df[col2], color=color2, marker='s', linestyle='--',
                     linewidth=2, markersize=6, label=col2)
    ax2.tick_params(axis='y', labelcolor=color2)
    ax2.legend(loc='upper right', fontsize=9)

    # 3. 关键修复：获取正确的x轴坐标和y轴范围
    # 3.1 将datetime索引转换为matplotlib内部数值（匹配x轴渲染）
    x_data = df.index  # 核心：datetime → 数值坐标
    # 3.2 获取调整后的y轴范围（绘制数据后再获取，确保覆盖所有数据）
    y_min, y_max = ax1.get_ylim()
    y_height = y_max - y_min  # 矩形高度（匹配左轴数据范围）

    same_trend = len(df['同趋势'])
    trend_1 = len(df[df['同趋势'] == 1])
    trend_0 = len(df[df['同趋势'] == 0])
    # 4. 绘制背景矩形（按时间区间填充）
    for i in range(1, len(df)):
        # 4.1 确定当前区间的x坐标范围（左边界=第i个日期，右边界=第i+1个日期）
        x_left = x_data[i-1]
        x_width = x_data[i] - x_left  # 区间宽度（按实际日期间隔）
        
        # 4.2 确定背景颜色（同趋势=1→粉红，=0→淡绿）
        bg_color = '#FFCCD5' if df['同趋势'].iloc[i] == 1 else '#CCFFCC'
        if df['GDP当季同比'].iloc[i] == df['GDP当季同比'].iloc[i-1]:
            bg_color = "#717171"
        # 4.3 添加矩形背景（zorder=-1确保在最底层，不遮挡线条）
        ax1.add_patch(Rectangle(
            xy=(x_left, y_min),    # 左下角坐标（x=当前日期，y=左轴最小值）
            width=x_width,         # 宽度=两个日期的间隔
            height=y_height,       # 高度=左轴数据范围
            color=bg_color,
            alpha=0.5,             # 透明度（避免遮挡数据）
            zorder=-1              # 图层优先级（-1最低，在所有数据之下）
        ))
        # print(i, str(x_left)[:10], df['同趋势'].iloc[i], bg_color)
    # 5. 图表美化
    title = title + f"（同趋势{trend_1}期，不同趋势{trend_0}期）"
    title = title + f"相同趋势率{100* trend_1/same_trend:.1f}%"
    plt.title(title, fontsize=12, pad=15)   
    ax1.grid(alpha=0.3, linestyle='-', linewidth=0.5)  # 网格线（淡色不抢镜）
    fig.tight_layout()  # 自动调整布局，避免标签被截断
    plt.close()
    return fig


def plot_heatmap(df, title='DataFrame热力图'):
    """
    绘制DataFrame热力图并返回figure对象
    
    参数:
    df: DataFrame，需要绘制热力图的数据
    title: 图表标题，默认为'DataFrame热力图'
    
    返回:
    fig: matplotlib的figure对象，可以用于进一步操作或保存
    """
    # 设置中文字体支持
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False    # 用来正常显示负号

    # 创建图形对象
    fig, ax = plt.subplots(figsize=(12, 8))
    
    # 绘制热力图
    sns.heatmap(df, 
                annot=True,      # 在单元格中显示数值
                fmt='.2f',       # 数值格式，保留两位小数
                cmap='RdBu_r',   # 颜色映射，红蓝渐变色
                center=0,        # 颜色中心点
                square=True,     # 单元格为正方形
                cbar_kws={'shrink': 0.8},  # 调整颜色条大小
                ax=ax)           # 指定绘图的axes

    ax.set_title(title)
    plt.tight_layout()
    
    # 返回figure对象而不是直接显示
    return fig

# 示例用法
if __name__ == "__main__":
    from _data_process import cumulative_to_monthly
    from _data_process import monthly_to_quarterly
    from _data_process import calculate_yoy

    series_data = pd.read_excel(f"data/全社会用电量当月累计值.xlsx", index_col=0)
    x_cum = series_data.index
    y_cum = series_data.iloc[:,0].values
    # 绘制折线图（宽度10，高度6）
    fig_cum = plot_line_chart(10, 6, 
                          x_cum, y_cum, 
                          "全社会用电量", "日期", "累计值")

    series_data_monthly = cumulative_to_monthly(series_data)
    data = series_data_monthly.iloc[:,0]
    x_monthly = data.index
    y_monthly = data.values
    
    # 绘制折线图（宽度10，高度6）
    fig_monthly = plot_line_chart(10, 6, 
                          x_monthly, y_monthly, "全社会用电量", "日期", "当月值")
    
    series_data_quarterly = monthly_to_quarterly(series_data_monthly)
    data = series_data_quarterly.iloc[:,0]
    x = data.index
    y = data.values
    # 绘制折线图（宽度10，高度6）
    fig_quarterly = plot_line_chart(10, 6, 
                          x, y, "全社会用电量", "日期", "季度值")


    series_data_quarterly_yoy = calculate_yoy(series_data_quarterly)
    data = series_data_quarterly_yoy.iloc[:,0]
    x_yoy = data.index
    y_yoy = data.values
    # 绘制折线图（宽度10，高度6）
    fig_quarterly_yoy = plot_line_chart(10, 6, 
                          x_yoy, y_yoy, "全社会用电量", "日期", "季度同比值")

    # 显示图像
    plt.show()
    # 如果需要保存图像，可以使用
    # fig.savefig("line_chart.png", bbox_inches='tight')